# -*- coding: UTF-8 -*-
"""
A brain of the agent && All decisions are made in here.
"""

import numpy as np
import pandas as pd


class RL(object):
    # learning_rate：学习率α
    # reward_decay：衰减度γ
    # e_greedy：epsilon_贪婪度
    # q_table:Q表
    # 纵轴：代表不同的state状态
    # 横轴：可选择的action行为
    def __init__(self, actions, learning_rate=0.01, reward_decay=0.9, e_greedy=0.9):
        self.actions = actions  # a list
        self.lr = learning_rate
        self.gamma = reward_decay
        self.epsilon = e_greedy
        self.q_table = pd.DataFrame(columns=self.actions)

    #  定义选择行为的方法
    def choose_action(self, observation):
        # 首先，检验当前state是否存在，如否，在Q表中加入此state
        self.check_state_exist(observation)
        # 选择 action
        if np.random.uniform() < self.epsilon:
            # 90%下 选择 Q value 最高的 action
            state_action = self.q_table.ix[observation, :]
            # 同一个 state, 可能会有多个相同的 Q action value,
            # 所以我们乱序一下，避免一直选取同一个索引操作
            state_action = state_action.reindex(np.random.permutation(state_action.index))
            action = state_action.argmax()
        else:  # 10%下 随机选择 action
            action = np.random.choice(self.actions)
        return action

    # 判断state是否存在过，如果第一次出现，将此信息放置Qtable中
    def check_state_exist(self, state):
        if state not in self.q_table.index:
            # 向q-table中增加新state
            self.q_table = self.q_table.append(
                pd.Series(
                    [0] * len(self.actions),
                    index=self.q_table.columns,
                    name=state,
                )
            )

    def learn(self, *args):
        pass


# off-policy
class QLearningTable(RL):
    def __init__(self, actions, learning_rate=0.01, reward_decay=0.9, e_greedy=0.9):
        super(QLearningTable, self).__init__(actions, learning_rate, reward_decay, e_greedy)

    # 定义Q学习方法，从现在的S中，和A以及R和Snext，来提升Qtable
    def learn(self, s, a, r, s_):
        # 检测 q_table 中是否存在 s_next
        self.check_state_exist(s_)
        # 计算Q估计
        q_predict = self.q_table.ix[s, a]
        if s_ != 'terminal':
            # 如果下个 state 不是 终止状态，计算Q真实
            q_target = r + self.gamma * self.q_table.ix[s_, :].max()
        else:
            # 下个 state 是终止符，直接获得奖励
            q_target = r
        # 更新对应的 state-action 值，学习率α*（Q真实-Q估计）
        self.q_table.ix[s, a] += self.lr * (q_target - q_predict)

# on-policy
class SarsaTable(RL):

    def __init__(self, actions, learning_rate=0.1, reward_decay=0.5, e_greedy=0.9):
        super(SarsaTable, self).__init__(actions, learning_rate, reward_decay, e_greedy)

    # 定义Sarsa学习方法，从现在的S中，和A以及R和Snext，来提升Qtable
    def learn(self, s, a, r, s_, a_):
        # 检测 q_table 中是否存在 s_next
        self.check_state_exist(s_)
        # get Q估计
        q_predict = self.q_table.ix[s, a]
        if s_ != 'terminal':
            # 如果下个 state 不是 终止状态，计算Q真实
            # 与Q学习区别就在于，Q学习的Q真实，是选取s_的最大值，而Sarsa是选了哪一步操作，就是哪一步：[s_, a_]
            q_target = r + self.gamma * self.q_table.ix[s_, a_]
        else:
            # 下个 state 是终止符，直接获得奖励
            q_target = r
        # 更新对应的 state-action 值，学习率α*（Q真实-Q估计）
        self.q_table.ix[s, a] += self.lr * (q_target - q_predict)  # update
